import os
import operator
from dotenv import load_dotenv
from typing import Annotated
# 新增: 导入 ToolMessage 用于手动创建工具响应
from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage
from langchain_core.tools import tool
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_openai import ChatOpenAI
from pydantic import Field, BaseModel
from langgraph.graph import StateGraph, END
# 移除: 不再需要预构建的 ToolNode
# from langgraph.prebuilt import ToolNode

# --- 1. 加载环境变量 ---
load_dotenv()

# --- 2. 定义工具 ---

# a. 一个总是能成功的工具
search_tool = DuckDuckGoSearchRun()

# b. 新增: 一个可能会失败的工具
@tool
def divide(a: float, b: float) -> float:
    """计算 a 除以 b 的结果。"""
    if b == 0:
        raise ValueError("错误：除数不能为零。")
    return a / b

# c. 将所有工具放入列表
tools = [search_tool, divide]

# --- 3. 定义状态 ---
class AgentState(BaseModel):
    messages: Annotated[list[BaseMessage], operator.add] = Field(default_factory=list)

# --- 4. 定义图中的节点 ---

# a. 初始化模型并绑定工具
llm = ChatOpenAI(model="qwen-plus-latest", base_url=os.getenv("OPENAI_BASE_URL"))
model_with_tools = llm.bind_tools(tools)

# b. Agent 节点 (无变化)
def agent_node(state: AgentState):
    print("---AGENT: 思考中...---")
    response = model_with_tools.invoke(state.messages)
    return {"messages": [response]}

# c. 自定义工具节点 (核心修改)
def custom_tool_node(state: AgentState):
    """
    这个自定义节点取代了预构建的 ToolNode。
    它包含错误处理逻辑。
    """
    print("---TOOLS: 执行工具中...---")
    # 获取 LLM 上一步请求调用的工具列表
    tool_calls = state.messages[-1].tool_calls
    
    # 将工具列表转换为 {name: tool} 的字典，方便查找
    tool_map = {tool.name: tool for tool in tools}
    
    # 准备一个列表来存放所有工具调用的返回结果
    tool_messages = []
    
    for call in tool_calls:
        tool_name = call["name"]
        tool_args = call["args"]
        
        print(f"---TOOLS: 准备调用 '{tool_name}'，参数: {tool_args}---")
        
        if tool_name not in tool_map:
            print(f"---TOOLS: 错误 - 未知的工具: {tool_name}---")
            # 如果模型试图调用一个不存在的工具，我们也需要通知它
            response_message = ToolMessage(
                content=f"Error: Tool '{tool_name}' not found.",
                tool_call_id=call["id"]
            )
        else:
            selected_tool = tool_map[tool_name]
            try:
                # 关键: 使用 try-except 块包裹工具的调用
                tool_output = selected_tool.invoke(tool_args)
                print(f"---TOOLS: '{tool_name}' 执行成功---")
                response_message = ToolMessage(
                    content=str(tool_output), # 确保内容是字符串
                    tool_call_id=call["id"]
                )
            except Exception as e:
                # 如果工具执行失败，捕获异常
                print(f"---TOOLS: '{tool_name}' 执行失败: {e}---")
                # 将错误信息格式化后存入 ToolMessage
                response_message = ToolMessage(
                    content=f"Error running tool '{tool_name}': {e}",
                    tool_call_id=call["id"]
                )
        
        tool_messages.append(response_message)
        
    return {"messages": tool_messages}

# --- 5. 定义条件边 (无变化) ---
def should_continue(state: AgentState) -> str:
    print("---ROUTER: 决策中...---")
    last_message = state.messages[-1]
    if not last_message.tool_calls:
        print("---ROUTER: 任务完成，结束流程。---")
        return "__end__"
    else:
        print("---ROUTER: 需要调用工具。---")
        return "tools"

# --- 6. 构建图 ---
workflow = StateGraph(AgentState)
workflow.add_node("agent", agent_node)
# 注册我们的自定义工具节点
workflow.add_node("tools", custom_tool_node)
workflow.set_entry_point("agent")
workflow.add_conditional_edges(
    "agent",
    should_continue,
)
workflow.add_edge("tools", "agent")
app = workflow.compile()

# --- 7. 运行图 ---
if __name__ == "__main__":
    # 准备一个会触发除零错误的问题
    question = "10除以0等于多少？"
    inputs = {"messages": [HumanMessage(content=question)]}
    
    print(f"开始执行任务: {question}\n")
    for event in app.stream(inputs, stream_mode="values"):
        event['messages'][-1].pretty_print()
        print("\n" + "="*50 + "\n")
        
    # 再试一个正常的问题
    question = "通义千问的最新版本是什么？"
    inputs = {"messages": [HumanMessage(content=question)]}
    
    print(f"\n开始执行新任务: {question}\n")
    for event in app.stream(inputs, stream_mode="values"):
        event['messages'][-1].pretty_print()
        print("\n" + "="*50 + "\n") 